package cc.shacocloud.luanniao.core.utils

import kotlinx.coroutines.reactor.awaitSingle
import kotlinx.coroutines.reactor.awaitSingleOrNull
import org.springframework.core.ResolvableType
import org.springframework.data.domain.Page
import org.springframework.data.domain.PageImpl
import org.springframework.data.domain.Pageable
import org.springframework.data.r2dbc.core.R2dbcEntityTemplate
import org.springframework.data.r2dbc.repository.R2dbcRepository
import org.springframework.data.relational.core.query.Criteria
import org.springframework.data.relational.core.query.CriteriaDefinition.Combinator
import org.springframework.data.relational.core.query.Query

/**
 *
 * @author 思追(shaco)
 */
object QueryExpand {

    private val r2dbcEntityTemplate: R2dbcEntityTemplate by lazy {
        ApplicationContextHolder.getInstance().getBean(R2dbcEntityTemplate::class.java)
    }

    /**
     * 使用 [Combinator.AND] 连接符号
     */
    fun and(vararg criteria: Criteria?): Criteria {
        val criteriaList = listOfNotNull(*criteria)
        return if (criteriaList.isEmpty()) Criteria.empty()
        else if (criteriaList.size == 1) criteriaList[0]
        else criteriaList[0].and(criteriaList.subList(1, criteriaList.size))
    }

    /**
     * 指定条件的数据是否存在
     */
    suspend fun <T, ID> R2dbcRepository<T, ID>.exists(criteria: Criteria): Boolean {
        val entityClass = getEntityClass()
        return r2dbcEntityTemplate.exists(Query.query(criteria), entityClass).awaitSingle()
    }

    /**
     * 查询一条数据
     */
    suspend fun <T, ID> R2dbcRepository<T, ID>.findOne(criteria: Criteria): T? {
        val entityClass = getEntityClass()
        return r2dbcEntityTemplate.selectOne(Query.query(criteria), entityClass).awaitSingleOrNull()
    }

    /**
     * 删除数据
     */
    suspend fun <T, ID> R2dbcRepository<T, ID>.deleteBy(criteria: Criteria): Long {
        val entityClass = getEntityClass()
        return r2dbcEntityTemplate.delete(Query.query(criteria), entityClass).awaitSingle()
    }

    /**
     * 分页查询
     */
    suspend fun <T, ID> R2dbcRepository<T, ID>.pageQuery(criteria: Criteria, pageable: Pageable): Page<T> {
        val entityClass = getEntityClass()

        // 不分页查询
        if (pageable.isUnpaged) {
            val contents = r2dbcEntityTemplate.select(entityClass)
                .matching(
                    Query.query(criteria).sort(pageable.sort)
                )
                .all()
                .collectList()
                .awaitSingle()

            return PageImpl(contents, pageable, contents.size.toLong())
        }

        // 查询total
        val total = r2dbcEntityTemplate.select(entityClass)
            .matching(
                Query.query(
                    criteria
                )
            )
            .count()
            .awaitSingle()

        if (total == 0.toLong()) {
            return PageImpl(emptyList(), pageable, total)
        }

        // 查询分页结果
        val contents = r2dbcEntityTemplate.select(entityClass)
            .matching(
                Query.query(criteria)
                    .sort(pageable.sort)
                    .offset(pageable.offset)
                    .limit(pageable.pageSize)
            )
            .all()
            .collectList()
            .awaitSingle()

        return PageImpl(contents, pageable, total)
    }

    /**
     * 获取存储库的实体类
     */
    @Suppress("UNCHECKED_CAST")
    fun <T, ID> R2dbcRepository<T, ID>.getEntityClass(): Class<T> {
        return ResolvableType.forClass(R2dbcRepository::class.java, javaClass).getGeneric(0)
            .toClass() as Class<T>
    }

}

